import torch
from torch import optim

from Expreriment.models.Expreriment import Expreriment

device = 'cuda' if torch.cuda.is_available() else 'cpu'


def config_model(args, dataset):
    if args.model == 'Expreriment':
        model = Expreriment(dataset=dataset, args=args)
    else:
        raise NotImplemented
    if args.model == 'Expreriment':
        optim_params = [{'params': model.parameters(), 'lr': args.lr}]
    else:
        raise NotImplemented
    model = model.to(device)
    optimizer = optim.Adam(optim_params, lr=args.lr, weight_decay=args.wd)
    return model, optimizer
